# ==============================================================================
# Randomly select one batch and derive sample data to train the network
# ==============================================================================

from scipy.io import savemat
import numpy as np
import random as rd
from filterpy.kalman import UnscentedKalmanFilter as UKF
from filterpy.kalman import JulierSigmaPoints as SP
from joblib import Parallel, delayed
# using the data with small range of distances (60->20)
from Trajectory_data_generator2 import *
threshold = 6
data_len = 50
# define UKF
dim_state = 4
DIM_STATE = 4
dim_observation = 2
DIM_OBSERVATION = 2
# transition noise
# Sampling time
sT = 0.1
state_n = 10.0
s_var = np.square(state_n)
T2 = np.power(sT, 2)
T3 = np.power(sT, 3)
T4 = np.power(sT, 4)
var_m = np.array([[T4 / 4, 0, T3 / 2, 0], [0, T4 / 4, 0, T3 / 2], [T3 / 2, 0, T2, 0], [0, T3 / 2, 0, T2]]) * s_var
# var_m = np.array([[T4/4,0,0,0],[0,T4/4,0,0],[0,0,T2,0],[0,0,0,T2]]) * s_var
# observation noise
dis_n = 10.0  # distance
DISTANCE_NOISE = 10.0
dis_var = np.square(dis_n)
DISTANCE_VARIANCE = np.square(DISTANCE_NOISE)
azi_n = 8.0  # azimuth
AZIMUTH_NOISE = 8.0
azi_var = np.square(azi_n / 1000)
AZIMUTH_VARIANCE = np.square(AZIMUTH_NOISE / 1000)


# Modify observations, make them continuous
# 处理具有周期性的数据（如角度）以消除因周期性跳跃导致的不连续性
def data_refine(observation_result):
    batch_size, data_len, _ = np.shape(observation_result)
    new_observation = np.copy(observation_result)
    for j in range(batch_size):
        for i in range(data_len - 1):
            a = new_observation[j, i, 0]
            b = new_observation[j, i + 1, 0]
            c = a - b
            if c > threshold:
                new_observation[j, i + 1:, 0] = new_observation[j, i + 1:, 0] + 2 * np.pi
            if c < -threshold:
                new_observation[j, i + 1:, 0] = new_observation[j, i + 1:, 0] - 2 * np.pi

    return new_observation


# Creat batch for training

# State transition function
def my_fx(x, sT):
    """ state transition function for state [downrange, vel, altitude]"""
    F_matrix = np.array([[1, 0, sT, 0], [0, 1, 0, sT], [0, 0, 1, 0], [0, 0, 0, 1]], 'float64')  # F_cv
#   F_matrix = np.array([[1, 0, np.sin(w*sT)/w, (np.cos(w*sT)-1)/w],
#                         [0, 1, (1-np.cos(w*sT))/w, np.sin(w*sT)/w],
#                         [0, 0, np.cos(w*sT), -np.sin(w*sT)],
#                         [0, 0, np.sin(w*sT), np.cos(w*sT)]],'float64')  # F_ct
    return np.dot(F_matrix, x)

def normalize_angle(angle):
    """将角度归一化到[-pi, pi]范围内"""
    while angle > np.pi:
        angle -= 2 * np.pi
    while angle < -np.pi:
        angle += 2 * np.pi
    return angle


def state_transition_ct(x, dt):
    """
    恒定转弯（CT）模型的状态转移函数
    参数：
    x: k-1 时刻的状态向量[px,py,v, yaw, omega]
    dt: 时间步长

    返回：
    k时刻的预测状态向量
    """
    px, py, v, yaw, omega = x  # 解包当前状态
    # 基于当前状态和dt预测下一状态 ---
    # 避免omega除零错误
    if abs(omega) < 1e-6:  # 判断接近零时转弯率的阈值
        # 近似线性运动（匀速直线运动）
        px_next = px + v * dt * np.cos(yaw)
        py_next = py + v * dt * np.sin(yaw)
    else:
        # 圆周运动方程
        v_over_omega = v / omega
        # 计算下一时刻的临时航向角，用于三角函数计算
        yaw_next_temp = yaw + omega * dt

        px_next = px + v_over_omega * (np.sin(yaw_next_temp) - np.sin(yaw))
        py_next = py + v_over_omega * (-np.cos(yaw_next_temp) + np.cos(yaw))

    # 更新航向角，速率和转弯率
    # 确定性模型部分假设速率和转弯率不变
    yaw_next = normalize_angle(yaw + omega * dt)  # 归一化最终的航向角
    v_next = v
    omega_next = omega

    return np.array([px_next, py_next, v_next, yaw_next, omega_next], dtype='float64')


# Observation function
def my_hx(x):
    """ returns slant range = np.array([[0],[0]]) based on downrange distance and altitude"""
    r = np.array([0, 0], 'float64')
    r[0] = np.arctan2(x[1], x[0])
    r[1] = np.sqrt(np.square(x[0]) + np.square(x[1]))
    return r


# Batch creating
def creat_batch3(pos_noise, vel_noise, BN):
    """
    :param pos_noise: 初始位置的噪声标准差
    :param vel_noise: 初始速度的噪声标准差
    :param BN: 控制批量大小的系数
    :return:
             trajectory_truth: 真实轨迹---(p_x,p_y,v_x,v_y)
             trajectory_observation: 观测轨迹---(theta,r)
             trajectory_estimate: UKF估计的轨迹---(p_x,p_y,v_x,v_y)
             trajectory_error: 真实轨迹与估计轨迹的误差（估计误差）
    """
    batch_size = int(100 * BN)
    # trajectory_batch_generator(batch_size, data_len)
    # return my_trajectory(batch_size, data_len,4):p_x,p_y,v_x,v_y
    # return my_observation(batch_size, data_len, 2):theta,r
    # return Tran_matrix(batch_size, 4,4):F_cv/F_ct
    my_trajectory, my_observation, tran_matrix = trajectory_batch_generator(batch_size, data_len)
    trajectory_truth = my_trajectory
    trajectory_observation = my_observation

    trajectory_estimate = np.zeros((batch_size, data_len, 4), dtype='float64')
    for i in range(batch_size):
        my_SP = SP(dim_state, kappa=0.)  # 生成sigma points
        my_UKF = UKF(dim_x=dim_state, dim_z=dim_observation, dt=sT, hx=my_hx, fx=my_fx, points=my_SP)
        my_UKF.Q *= var_m   # process noise covariance
        my_UKF.R *= np.array([[azi_var, 0], [0, dis_var]])  # observation noise covariance
        x0_noise = np.array([np.random.normal(0, pos_noise, 2), np.random.normal(0, vel_noise, 2)])
        my_UKF.x = trajectory_truth[i, 0, :] + np.reshape(x0_noise, [4, ])  # initial state + noise
        my_UKF.P *= 1.0  # initial state covariance

        # tracking results of UKF
        xs = []
        xs.append(my_UKF.x)
        for j in range(data_len - 1):
            my_UKF.predict()    # predict state
            my_UKF.update(trajectory_observation[i, j + 1, :])   # update state with measurement
            xs.append(my_UKF.x.copy())  # save state
        trajectory_estimate[i] = np.asarray(xs)

    return trajectory_truth, trajectory_observation, trajectory_estimate, trajectory_truth - trajectory_estimate


def create_all_data(batch_size, data_len, ini_pos_noise, ini_vel_noise, batch_num):
    """
    :param batch_size:批大小，即一个批次包含多少条轨迹，一般和model.fit函数中的batch_size保持一致，每个批次的turn_rate应该一致，常见取值为32,64
    :param data_len: 轨迹长度，即每条轨迹包含多少个状态点，也可以理解为另一种名称，即时间步数time_step_size，常见取值为50,100等
    :param ini_pos_noise:初始位置的噪声标准差，
    :param ini_vel_noise:初始速度的噪声标准差
    :param batch_num:批次数量，即生成多少批次的数据,即每个epoch或每update_frequency个batch更新一次权重，常见取值为100，200等
    :return:all_true_orbit:所有真实轨迹---(batch_size*batch_num, data_len, 4)，其中状态量顺序为[p_x,p_y,v_x,v_y]
             all_observation_orbit:所有观测轨迹---(batch_size*batch_num, data_len, 2),其中观测量顺序为[theta,r]
             all_estimate_orbit:所有由UKF估计的轨迹---(batch_size*batch_num, data_len, 4)，其中UKF估计的状态量顺序为[p_x,p+y,v_x,v_y]
             all_error_between:all_true_orbit与all_estimate_orbit之间的误差（前者减去后者）---(batch_size*batch_num, data_len, 4),其中的值可能为负值
    """
    all_true_orbit = np.zeros((batch_size*batch_num, data_len, 4), dtype='float32')
    all_observation_orbit = np.zeros((batch_size * batch_num, data_len, 2), dtype='float32')
    all_estimate_orbit = np.zeros((batch_size * batch_num, data_len, 4), dtype='float32')

    # 并行生成所有批次的轨迹数据
    def process_batch(batch_idx):
        true_orbit, observation_orbit, _ = trajectory_batch_generator(batch_size, data_len)
        estimate_orbit = np.zeros((batch_size, data_len, 4), dtype='float32')  # 用于保存UKF估计的轨迹

        # 为每个批次并行处理UKF
        for j in range(batch_size):
            my_sp = SP(DIM_STATE, kappa=0.)  # 生成sigma points
            my_ukf = UKF(dim_x=DIM_STATE, dim_z=DIM_OBSERVATION, dt=sT, hx=my_hx, fx=my_fx, points=my_sp)
            my_ukf.Q *= var_m  # process noise covariance
            my_ukf.R *= np.array([[AZIMUTH_VARIANCE, 0], [0, DISTANCE_VARIANCE]])  # 观测噪声协方差
            x0_noise = np.array([np.random.normal(0, ini_pos_noise, 2), np.random.normal(0, ini_vel_noise, 2)])
            my_ukf.x = true_orbit[j, 0, :] + np.reshape(x0_noise, [4, ])  # 初始状态+噪声
            my_ukf.P *= 1.0  # 初始化状态协方差

            # 跟踪UKF估计结果
            xs = []
            xs.append(my_ukf.x)
            for k in range(data_len-1):
                my_ukf.predict()  # 预测状态
                my_ukf.update(observation_orbit[j, k+1, :])
                xs.append(my_ukf.x.copy())  # 保存状态
            estimate_orbit[j] = np.asarray(xs)  # 保存UKF估计的轨迹

        if batch_idx != 0 and (batch_idx+1) % 50 == 0:
            print("生成第 %d 批数据" % (batch_idx + 1))

        return true_orbit, observation_orbit, estimate_orbit

    # 使用joblib并行处理
    results = Parallel(n_jobs=-1)(delayed(process_batch)(i) for i in range(batch_num))

    # 组装结果
    for i, (true_orbit, observation_orbit, estimate_orbit) in enumerate(results):
        all_true_orbit[i * batch_size:(i + 1) * batch_size] = true_orbit
        all_observation_orbit[i * batch_size:(i + 1) * batch_size] = observation_orbit
        all_estimate_orbit[i * batch_size:(i + 1) * batch_size] = estimate_orbit

    all_error_orbit = all_true_orbit - all_estimate_orbit
    return all_true_orbit, all_observation_orbit, all_estimate_orbit, all_error_orbit


def create_all_data_optimized(batch_size, data_len, ini_pos_noise, ini_vel_noise, batch_num):
    """
    使用joblib并行处理生成轨迹数据
    :param batch_size: 批大小
    :param data_len: 轨迹长度
    :param ini_pos_noise: 初始位置噪声标准差
    :param ini_vel_noise: 初始速度噪声标准差
    :param batch_num: 批次数量
    :return: 所有轨迹数据
    """
    all_true_orbit =

if __name__ == '__main__':
    all_true_orbit, all_observation_orbit, all_estimate_orbit, all_error_orbit = create_all_data(batch_size=32, data_len=50, ini_pos_noise=30, ini_vel_noise=3, batch_num=100)
    #trajectory_truth, trajectory_observation, trajectory_estimate, error_between = creat_batch3(pos_noise=0.01, vel_noise=0., BN=0.1)
    filename = './matlab/pre_process.mat'
    savemat(filename, {'all_true_orbit': all_true_orbit, 'all_observation_orbit': all_observation_orbit, 'all_estimate_orbit': all_estimate_orbit, 'all_error_orbit': all_error_orbit})